package org.ala.tree.bt;

import java.util.Comparator;
import java.util.stream.IntStream;

import org.ala.tree.AbstractTree;

/**
 * B树
 * <p>	自平衡树（处在同一层节点的高度一致）
 * <p>	对于m阶b树，每个节点存储m-1个的关键字，同时对下层持有m的子节点
 * <p>		比如当前节点[5,10,15]，其4个子节点表示的范围依次是(-∞,5],(5,10],(10, 15],(15,+∞)
 * <p>	搜索：从根开始，依次找到子节点，最后找到叶子节点（没有子节点的节点）
 * <p>	插入：从根节点开始，找到叶子节点，把数据放入叶子结点
 * <p>		如果叶子节点的关键字 >= 阶数，则该节点要分裂
 * <p>	分裂：取本层节点数 / 2 上取整，取到分裂关键字
 * <p>		当前节点从分裂关键字的左右两边，分裂成2个节点
 * <p>		将分裂关键字保存进父节点中，在父节点中，新增关键字的左右两边分别指向新分裂出的2个节点
 * <p>		如果父节点关键字数 >= 阶数，继续往上分裂
 *
 * @author ala
 * @date 2025年5月20日
 */
@SuppressWarnings("unchecked")
public class BTree<K, V> extends AbstractTree<K, V> {

	
	/**
	 * 阶数
	 */
	protected int m;
	/**
	 * 	每层最小节点数
	 */
	protected int mx;
	/**
	 * 根节点
	 */
	protected Node<K,V> HEAD;
	
	public BTree(int m, Comparator<K> comparator) {
		super(comparator);
		this.m = m;
		this.mx = m >> 1;
	}

	@Override
	public boolean add(K k, V v) {
		if (size == 0) {
			HEAD = new Node<>(m);
			HEAD.keys[HEAD.kp++] = new Data<>(k, v);
			HEAD.lear = true;
			size++;
			return true;
		}
		if (_add(new Data<>(k,v))) {size++; return true;}
		return false;
	}
	/**
	 * 	找到数据的插入位置，执行插入，必要时分裂节点
	 * 	<p>	不允许key重复
	 */
	protected boolean _add(Data<K,V> d) {
		//	找到待插入的叶子节点
		SearchResult<K, V> sr = _search(HEAD, -1, d.key);
		if (sr.find) { return false; }
		
		//	此时找到的必为叶子节点
		Node<K,V> n = sr.node;
		n.keys[n.kp++] = d;
		
		//	如果kp到达了m-1，执行分裂
		if (n.kp == m) {
			_split(n);
		}
		return true;
	}
	protected void _split(Node<K,V> n) {
		//	[0, i)保留，(i,kp)生成新节点
		int i = n.kp >> 1;
		Data<K,V> d = n.keys[i];
		//	生成新节点，并把(i,kp)复制过去。并且nodes的parent指向nn
		Node<K,V> nn = new Node<>(m);
		nn.lear = n.lear;
		for (int j = i+1 ; j < n.kp ; j++) {
			nn.keys[nn.kp] = n.keys[j];
			if (!n.lear) {
				nn.nodes[nn.kp] = n.nodes[j];
				n.nodes[j].parent = nn;
			}
			nn.kp++;
		}
		nn.nodes[nn.kp] = n.nodes[n.kp];
		if (!n.lear) { n.nodes[n.kp].parent = nn; }
		n.kp = i;
		n.nodes[n.kp] = n.nodes[i];
		
		
		//	i节点上升到父节点
		Node<K,V> p = n.parent;
		if (p == null) {
			p = new Node<>(m);
			HEAD = p;
			n.parent = p;
		}
		nn.parent = p;
		//	把i插入到父节点
		int j = p.kp;
		for (int t = 0 ; t < p.kp ; t++) {
			int c = comparator.compare(p.keys[t].key, d.key);
			if (c >= 0) {j = t; break;}
		}
		if (j == p.kp) {
			p.keys[j] = d;
			p.nodes[j] = n;
			p.nodes[j+1] = nn;
		} else {
			for (int t = p.kp ; t > j ; t--) {
				p.keys[t] = p.keys[t-1];
			}
			p.keys[j] = d;
			p.nodes[j] = n;
			p.nodes[j+1] = nn;
		}
		p.kp++;
		
		//	检测父节点如果也溢出了，递归分裂父节点
		if (p.kp == m) {
			_split(p);
		}
	}

	@Override
	public boolean remove(K k) {
		//	找到要删除的节点
		SearchResult<K, V> sr = _search(k);
		if (!sr.find) {return false;}
		
		Node<K,V> n = sr.node; int ni = sr.ni;
		int di = sr.di;
		//	删除第i个节点
		_remove(n, ni, di);
		return true;
	}
	/**
	 * 	在node中删除第i个节点
	 */
	protected void _remove(Node<K,V> n, int ni, int i) {
		//	如果是叶子节点，直接删
		if (n.lear) {
			for (int j = i ; j < n.kp-1 ; j++) {
				n.keys[j] = n.keys[j+1];
			}
			n.kp--;
		}
		//	如果删除的是非叶节点，把中序遍历的后继替换到要删除的节点。接着删除中序遍历后继即可
		else {
			//	oracle在删数据时，不可能全量遍历，再拿到中序后继
		}
		//	判断是否出现下溢
		tryMerge(n, ni);
	}
	/**
	 * 	合并节点
	 *  1）检测左右兄弟，是否够借，若够借：
	 *  	将父节点下移到当前节点
	 *  	将左右兄弟的临界节点上移到父节点
	 *  2）若都不够借
	 *  	把父节点中的关键字下移到本节点中，再合并左兄弟
	 *  	父节点中删除被下移的关键字
	 */
	protected void tryMerge(Node<K,V> n, int ni) {
		//	如果n没有发生下溢，直接返回
		if (n.kp < (m >> 1)) {
			return;
		}
		//	如果是根节点（根节点的ni是-1），没得借。那没辙，删了白删
		if (n.parent == null) {
			if (n.kp == 0) {
				HEAD = null;
			}
		}
		//	检测左右兄弟，是否够借
		else {
			Node<K,V> p = n.parent, nl = ni > 0 ? p.nodes[ni-1] : null, nr = ni < p.kp-1 ? p.nodes[ni+1] : null;
			//	左右兄弟不可能全为空
			if (nl != null && nr != null) {
				//	优先找左兄弟借
				if (nl.kp >= mx) {
					borrow(p, n, ni, nl, true);
				} 
				//	左兄弟没余粮，找有兄弟借
				else if (nr.kp >= mx) {
					borrow(p, n, ni, nr, false);
				}
				//	左右都没余粮，执行合并
				else {
					//	哪边小往哪边合
					if (nl.kp <= nr.kp) {
						merge(p, n, ni, nl, true);
						if (p.parent == null) {ni = 0;}
						else {
							for (int i = 0 ; i <= p.kp ; i++) {
								if (p.parent.nodes[i] == p) {ni = i; break;}
							}
						}
						tryMerge(p, ni);
					} else {
						merge(p, n, ni, nr, false);
						if (p.parent == null) {ni = 0;}
						else {
							for (int i = 0 ; i <= p.kp ; i++) {
								if (p.parent.nodes[i] == p) {ni = i; break;}
							}
						}
						tryMerge(p, ni);
					}
				}
			}
			//	左边不为空
			else if (nl != null) {
				if (nl.kp >= mx) { borrow(p, n, ni, nl, true); }
				else {
					merge(p, n, ni, nl, true);
					if (p.parent == null) {ni = 0;}
					else {
						for (int i = 0 ; i <= p.kp ; i++) {
							if (p.parent.nodes[i] == p) {ni = i; break;}
						}
					}
					tryMerge(p, ni);
				}
			} else {
				if (nr.kp >= mx) { borrow(p, n, ni, nr, false); }
				else {
					merge(p, n, ni, nr, false);
					if (p.parent == null) {ni = 0;}
					else {
						for (int i = 0 ; i <= p.kp ; i++) {
							if (p.parent.nodes[i] == p) {ni = i; break;}
						}
					}
					tryMerge(p, ni);
				}
			}
		}
	}
	/**
	 * 	找兄弟借一个节点
	 * 	<p>	只改关键字，需指针修改
	 *  <p>	因为父节点的节点数并没有减少，所以不会触发父节点的再均衡
	 *  <p>	所以这种情况只会发生在叶子节点
	 *  @param	nb brother节点（不是NB节点）
	 */
	protected void borrow(Node<K,V> p, Node<K,V> n, int ni, Node<K,V> nb, boolean left) {
		//	父节点下移到n
		Data<K,V> d = p.keys[ni];
		//	找到下放节点的插入位置
		int j = n.kp;
		for (int t = 0 ; t < n.kp ; t++) {
			if (comparator.compare(d.key, n.keys[t].key) >= 0) {
				j = t;
				break;
			}
		}
		//	执行插入
		if (j == n.kp) {
			n.keys[j] = d;
		} else {
			for (int t = n.kp ; t > j ; t--) {
				n.keys[t] = n.keys[t-1];
			}
			n.keys[j] = d;
		}
		//	兄弟的临界节点上移到父
		if (left) {
			p.keys[ni] = nb.keys[nb.kp-1];
			nb.kp--;
		} else {
			p.keys[ni] = nb.keys[0];
			for (int i = 0 ; i < nb.kp-1 ; i++) {
				nb.keys[i] = nb.keys[i+1];
			}
		}
	}
	/**
	 * 	合并
	 *  <p>	往左合：
	 *  		父节点ni的key下移到nl的末尾
	 *  		n追加到nl的末尾
	 *  		父节点keys，ni的key删掉
	 *  		父节点nodes，从ni+1位置开始整体左移（ni+1原来是指向n的，现在n无了）
	 *  <p>	往右合：
	 *  		父节点ni的key下移到n的末尾
	 *  		n追加到nr的开头
	 *  		父节点keys，ni的key删掉
	 *  		父节点nodes，从ni+1位置开始整体左移（ni+1原来是指向n的，现在n无了）
	 */
	protected void merge(Node<K,V> p, Node<K,V> n, int ni, Node<K,V> nb, boolean left) {
		Data<K,V> d = p.keys[ni];
		if (left) {
			//	父节点ni的key下移到nl的末尾
			nb.keys[nb.kp++] = d;
			//	n追加到nl的末尾
			for (int i = 0 ; i < n.kp ; i++) {
				nb.keys[nb.kp] = n.keys[i];
				nb.nodes[nb.kp++] = n.nodes[i];
			}
			nb.nodes[nb.kp] = n.nodes[n.kp];
			//	父节点keys，ni的key删掉
			//	父节点nodes，从ni+1位置开始整体左移（ni+1原来是指向n的，现在n无了）
			for (int i = ni ; i < p.kp-1 ; i++) {
				p.keys[i] = p.keys[i+1];
				p.nodes[i+1] = p.nodes[i+2];
			}
		} else {
			//	父节点ni的key下移到n的末尾
			n.keys[n.kp++] = d;
			//	n追加到nr的开头
			nb.nodes[nb.kp + n.kp] = nb.nodes[nb.kp];
			for (int i = nb.kp + n.kp - 1 ; i >= n.kp ; i--) {
				nb.keys[i] = nb.keys[i - n.kp];
				nb.nodes[i] = nb.nodes[i - n.kp];
			}
			//	原n的末尾，现在是下放的d的左边
//			nb.nodes[n.kp] = n.nodes[n.kp];
			for (int i = 0 ; i < n.kp ; i++) {
				nb.keys[i] = n.keys[i];
				nb.nodes[i] = n.nodes[i];
			}
			//	父节点keys，ni的key删掉
			//	父节点nodes，从ni+1位置开始整体左移（ni+1原来是指向n的，现在n无了）
			for (int i = ni ; i < p.kp-1 ; i++) {
				p.keys[i] = p.keys[i+1];
				p.nodes[i+1] = p.nodes[i+2];
			}
		}
	}
	

	@Override
	public boolean update(K k, V v) {
		SearchResult<K,V> sr = _search(HEAD, -1, k);
		if (sr.find) {
			sr.data.val = v;
			return true;
		}
		return false;
	}

	@Override
	public V search(K k) {
		SearchResult<K,V> sr = _search(HEAD, -1, k);
		return sr.find ? sr.data.val : null;
	}
	/**
	 * 	找k
	 * 	<p>	如果找不到，返回k需要插入的叶子节点
	 */
	protected SearchResult<K,V> _search(K k) {
		Node<K,V> crt = HEAD;
		int di = -1;
		while (!crt.lear) {
			//	阶数很高的情况下可以优化成二分
			for (int i = 0 ; i < crt.kp ; i++) {
				int c = comparator.compare(k, crt.keys[i].key);
				if (c < 0) {
					crt = crt.nodes[i];
					di = i;
				} else if (c == 0) {
					return SearchResult.build(crt, di, crt.keys[i], i);
				}
			}
			crt = crt.nodes[crt.kp];
			di = crt.kp;
		}
		return SearchResult.build(crt, di);
	}
	protected SearchResult<K,V> _search(Node<K,V> r, int ni, K k){
		for (int i = 0 ; i < r.kp ; i++) {
			int c = comparator.compare(k, r.keys[i].key);
			if (c == 0) {
				return SearchResult.build(r, ni, r.keys[i], i);
			} else if (c < 0) {
				if (r.lear) {
					return SearchResult.build(r, ni);
				}
				return _search(r.nodes[i], i, k);
			}
		}
		if (r.lear) {
			return SearchResult.build(r, ni);
		}
		return _search(r.nodes[r.kp], r.kp, k);
	}

	
	@Override
	public String toString() {
		StringBuilder sbuf = new StringBuilder();
		toString(sbuf, HEAD, 0);
		return sbuf.toString();
	}
	protected void toString(StringBuilder sbuf, Node<K,V> n, int level) {
		if (n == null) {return;}
		StringBuilder t = new StringBuilder();
		IntStream.range(0, level).forEach(i -> t.append("\t"));
		for (int i = 0 ; i < n.kp ; i++) {
			sbuf.append(t).append(n.keys[i].key).append("\n");
			toString(sbuf, n.nodes[i], level+1);
		}
		toString(sbuf, n.nodes[n.kp], level+1);
	}


	/**
	 * 	节点信息
	 */
	public static class Node<K,V> {
		/**
		 * 	父节点
		 */
		Node<K,V> parent;
		/**
		 * 	关键字数组（长为m）
		 * 	<p>	一旦长度到达m，就要分裂
		 */
		Data<K,V>[] keys;
		/**
		 * 	keys下一个可写指针
		 */
		int kp = 0;
		/**
		 * 	关键字对应的下级节点（长为m）
		 * 	<p>	keys[i]对应nodes[i]
		 * 	<p>	nodes[m-1]对应 > keys[m-2]
		 */
		Node<K,V>[] nodes;
		/**
		 * 	是否叶节点
		 */
		boolean lear;
		public Node(int m) {
			this.keys = new Data[m];
			this.nodes = new Node[m+1];		//	第m个节点存放 > keys[m]的指针
		}
		@Override
		public String toString() {
			StringBuilder sbuf = new StringBuilder();
			sbuf.append("[");
			for (int i = 0 ; i < kp ; i++) {
				if (i > 0) {sbuf.append(",");}
				sbuf.append(keys[i]);
			}
			sbuf.append("]");
			return sbuf.toString();
		}
	}
	/**
	 * 	节点数据
	 */
	public static class Data<K,V> {
		K key;
		V val;
		public Data(K k, V v) {key = k; val = v;}
		public String toString() {
			return String.format("(%s,%s)", key, val);
		}
	}
	/**
	 * 	搜索结果
	 */
	public static class SearchResult<K,V> {
		boolean find;
		Data<K,V> data;
		/**
		 * 	data在node中的下标
		 */
		int di;
		Node<K,V> node;
		/**
		 * 	node在parent中的下标
		 */
		int ni;
		public static <K,V> SearchResult<K,V> build(Node<K,V> node, int ni, Data<K,V> data, int di) {
			SearchResult<K,V> r = new SearchResult<>();
			r.find = true;
			r.data = data;
			r.di = di;
			r.node = node;
			r.ni = ni;
			return r;
		}
		public static <K,V> SearchResult<K,V> build(Node<K,V> node, int ni) {
			SearchResult<K,V> r = new SearchResult<>();
			r.find = false;
			r.data = null;
			r.di = -1;
			r.node = node;
			r.ni = ni;
			return r;
		}
	}
	
}
